package sentinel

import (
	"context"
	sentinel "github.com/alibaba/sentinel-golang/api"
	"github.com/alibaba/sentinel-golang/core/base"
	"github.com/alibaba/sentinel-golang/core/circuitbreaker"
	"github.com/alibaba/sentinel-golang/core/config"
	"github.com/alibaba/sentinel-golang/core/flow"
	"github.com/asim/go-micro/v3/client"
	"log"
)

type clientWrapper struct {
	client.Client
}

func InitSentinel() {
	conf := config.NewDefaultConfig()
	err := sentinel.InitWithConfig(conf)
	if err != nil {
		log.Fatal(err)
	}
	_, err = flow.LoadRules([]*flow.Rule{
		// 定义某个规则的限流策略
		{
			Resource:               "gateway",
			TokenCalculateStrategy: flow.Direct,
			ControlBehavior:        flow.Reject,
			Threshold:              10,
			StatIntervalInMs:       1000,
		},
	})
	if err != nil {
		log.Fatalf("Unexpected error: %+v", err)
		return
	}
	// 熔断策略
	_, err = circuitbreaker.LoadRules([]*circuitbreaker.Rule{
		// 慢调用比例规则
		{
			Resource:         "gateway",
			Strategy:         circuitbreaker.SlowRequestRatio,
			RetryTimeoutMs:   5000,
			MinRequestAmount: 10,
			StatIntervalMs:   10000,
			MaxAllowedRtMs:   20,
			Threshold:        0.1,
		},
		// 错误比例规则
		{
			Resource:         "gateway",
			Strategy:         circuitbreaker.ErrorRatio,
			RetryTimeoutMs:   5000,
			MinRequestAmount: 10,
			StatIntervalMs:   10000,
			Threshold:        0.1,
		},
		// 错误计数规则
		{
			Resource:         "gateway",
			Strategy:         circuitbreaker.ErrorCount,
			RetryTimeoutMs:   5000,
			MinRequestAmount: 10,
			StatIntervalMs:   10000,
			Threshold:        100,
		},
	})
	if err != nil {
		log.Fatalf("Unexpected error: %+v", err)
		return
	}
}

// Call 配置入站的Web请求
func (c *clientWrapper) Call(ctx context.Context, req client.Request, rsp interface{}, opts ...client.CallOption) error {

	entry, blockErr := sentinel.Entry(
		"gateway",
		sentinel.WithResourceType(base.ResTypeWeb),
		sentinel.WithTrafficType(base.Inbound),
	)

	if blockErr != nil {
		return blockErr
	}
	defer entry.Exit()

	err := c.Client.Call(ctx, req, rsp, opts...)
	if err != nil {
		sentinel.TraceError(entry, err)
	}

	return err
}

func (c *clientWrapper) Stream(ctx context.Context, req client.Request, opts ...client.CallOption) (client.Stream, error) {
	//resourceName := req.Method()

	entry, blockErr := sentinel.Entry(
		"gateway",
		sentinel.WithResourceType(base.ResTypeRPC),
		sentinel.WithTrafficType(base.Outbound),
	)

	if blockErr != nil {

		return nil, blockErr
	}
	defer entry.Exit()

	stream, err := c.Client.Stream(ctx, req, opts...)
	if err != nil {
		sentinel.TraceError(entry, err)
	}

	return stream, err
}

// NewWebRateLimitWrapper returns a sentinel client Wrapper.
func NewWebRateLimitWrapper() client.Wrapper {
	return func(c client.Client) client.Client {
		return &clientWrapper{c}
	}
}
